# -*- coding: utf-8 -*-
"""
# --------------------------------------------------------
# @Author : Pan
# @E-mail : 
# @Date   : 2025-09-29 08:51:42
# @Brief  : 视频编码演示脚本 - 用于多模态模型的视频预处理
#           主要功能：视频帧采样、时间戳映射、帧分组打包
# --------------------------------------------------------
"""

import torch
from PIL import Image
from modelscope import AutoModel, AutoTokenizer
from decord import VideoReader, cpu  # pip install decord
from scipy.spatial import cKDTree
import numpy as np
import math
from pybaseutils import image_utils

# 全局常量定义
MAX_NUM_FRAMES = 180  # 视频打包后接收的最大帧数，实际有效帧数为 MAX_NUM_FRAMES * MAX_NUM_PACKING
MAX_NUM_PACKING = 3   # 视频帧的最大打包数量，有效范围：1-6
TIME_SCALE = 0.1      # 时间尺度，用于时间戳的离散化处理


def group_array(arr, size):
    """
    将数组按指定大小分组
    
    Args:
        arr: 输入数组
        size: 每组的大小
    
    Returns:
        list: 分组后的数组列表
    """
    return [arr[i:i + size] for i in range(0, len(arr), size)]


def map_to_nearest_scale(values, scale):
    """
    使用KD树将值映射到最近的尺度点
    
    Args:
        values: 需要映射的值数组
        scale: 尺度数组（参考点）
    
    Returns:
        numpy.ndarray: 映射到最近尺度点的值数组
    """
    # 构建KD树用于快速最近邻搜索
    tree = cKDTree(np.asarray(scale)[:, None])
    # 查询每个值最近的尺度点
    _, indices = tree.query(np.asarray(values)[:, None])
    # 返回对应的尺度值
    return np.asarray(scale)[indices]


def encode_video(video_path, choose_fps=3, force_packing=None):
    """
    视频编码主函数 - 对视频进行帧采样和时间戳处理
    
    Args:
        video_path (str): 视频文件路径
        choose_fps (int, optional): 目标采样帧率，默认3fps
        force_packing (int, optional): 强制打包数量，None表示自动计算
    
    Returns:
        tuple: (frames, frame_ts_id_group)
            - frames: PIL图像列表
            - frame_ts_id_group: 分组后的时间戳ID列表
    """
    
    def uniform_sample(l, n):
        """
        均匀采样函数 - 从列表l中均匀采样n个元素
        
        Args:
            l: 输入列表
            n: 采样数量
        
        Returns:
            list: 采样结果列表
        """
        # 计算采样间隔
        gap = len(l) / n
        # 计算采样索引（居中采样）
        idxs = [int(i * gap + gap / 2) for i in range(n)]
        return [l[i] for i in idxs]

    # 使用decord读取视频
    vr = VideoReader(video_path, ctx=cpu(0))
    # 获取视频的原始帧率
    fps = vr.get_avg_fps()
    # 计算视频总时长（秒）
    video_duration = len(vr) / fps

    # 动态计算打包策略
    # 如果目标帧率下的总帧数不超过最大帧数限制，则不需要打包
    if choose_fps * int(video_duration) <= MAX_NUM_FRAMES:
        packing_nums = 1  # 单帧打包
        # 计算实际采样帧数，考虑帧率限制和最大帧数限制
        choose_frames = round(min(choose_fps, round(fps)) * min(MAX_NUM_FRAMES, video_duration))
    else:
        # 需要打包处理的情况
        # 计算需要的打包数量
        packing_nums = math.ceil(video_duration * choose_fps / MAX_NUM_FRAMES)
        if packing_nums <= MAX_NUM_PACKING:
            # 如果打包数量在允许范围内，按目标帧率采样
            choose_frames = round(video_duration * choose_fps)
        else:
            # 如果超过最大打包限制，使用最大允许的帧数
            choose_frames = round(MAX_NUM_FRAMES * MAX_NUM_PACKING)
            packing_nums = MAX_NUM_PACKING

    # 生成所有帧的索引
    frame_idx = [i for i in range(0, len(vr))]
    # 均匀采样得到目标帧索引
    frame_idx = np.array(uniform_sample(frame_idx, choose_frames))

    # 如果指定了强制打包数量，使用指定的值
    if force_packing:
        packing_nums = min(force_packing, MAX_NUM_PACKING)

    # 输出视频信息
    print(video_path, ' duration:', video_duration)
    print(f'get video frames={len(frame_idx)}, packing_nums={packing_nums}')

    # 批量获取帧数据并转换为numpy数组
    frames = vr.get_batch(frame_idx).asnumpy()

    # 计算每个帧的时间戳（秒）
    frame_idx_ts = frame_idx / fps
    # 生成时间尺度数组，用于时间戳离散化
    scale = np.arange(0, video_duration, TIME_SCALE)

    # 将时间戳映射到最近的尺度点，并转换为时间戳ID
    frame_ts_id = map_to_nearest_scale(frame_idx_ts, scale) / TIME_SCALE
    frame_ts_id = frame_ts_id.astype(np.int32)

    # 确保帧数和时间戳数量一致
    assert len(frames) == len(frame_ts_id)

    # 将numpy数组转换为PIL图像（RGB格式）
    frames = [Image.fromarray(v.astype('uint8')).convert('RGB') for v in frames]
    # 按打包数量对时间戳ID进行分组
    frame_ts_id_group = group_array(frame_ts_id, packing_nums)

    return frames, frame_ts_id_group


# 主程序入口
if __name__ == "__main__":
    # 视频文件路径
    video_path = "../../data/video2.mp4"
    # 目标采样帧率
    fps = 5  # fps for video
    # 强制打包设置，None表示自动计算
    force_packing = None  # You can set force_packing to ensure that 3D packing is forcibly enabled; otherwise, encode_video will dynamically set the packing quantity based on the duration.
    
    # 调用视频编码函数
    frames, frame_ts_id_group = encode_video(video_path, fps, force_packing=force_packing)
    
    # 显示每一帧图像（用于调试）
    for rgb in frames:
        image = np.asarray(rgb)
        image_utils.show_image_plt("image", image, use_rgb=True)

    # 输出总帧数
    print(len(frames))